import datetime
import logging
import time
import traceback

import kafka

from collections import defaultdict
from coin.base.datetime_util import to_timestamp_int
from coin.proto.coin_order_gateway_pb2 import (
    OrderGatewayLog, OrderEvent)
from coin.proto.coin_query_pb2 import (
    CurrencyTransferBundle,
    ProductFillElement,
    ProductFillBundle)
from coin.tool.strat_monitor.util.monitor_util import (
    _filter_account_topic)
from coin.util.queue.config import KafkaConfig
from coin.util.queue.tools.kafka_archive import run_from_kafka_archive

from xunkemgmt_client.support.accounting.database.db_output import (
    DBOutputBalance, DBOutputPosition, DBOutputTransfer, DBOutputAccountFill)
from xunkemgmt_client.support.accounting.database.db_querier import (
    AccountInfoQuerier)
from xunkemgmt_client.client.util.query_util import query_accounts


def _validate_account_request(account_request):
  assert account_request.HasField('market_type')
  assert account_request.HasField('exchange')
  assert account_request.HasField('owner')


class KafkaAccountInfoQuerier(AccountInfoQuerier):
  def __init__(self, kafka_config):
    self._kafka_config = kafka_config
    self._topics = kafka.KafkaConsumer(
        bootstrap_servers=self._kafka_config.kafka_servers).topics()
    self._balance = {}
    self._position = {}
    self._transfer = defaultdict(list)
    self._fill = defaultdict(dict)
    self._acct_key_func = lambda x: (x.market_type, x.exchange, x.owner)
    self._fill_key_func = lambda x: (x.symbol, x.external_order_id, x.fill_id)

  def _on_log(self, partition, record):
    pb = OrderGatewayLog()
    pb.ParseFromString(record.value)
    assert pb.HasField('account_request')
    _validate_account_request(pb.account_request)
    acct_key = self._acct_key_func(pb.account_request)
    if pb.type == OrderGatewayLog.BALANCE:
      assert pb.HasField('balance')
      self._balance[acct_key] = pb.balance
    elif pb.type == OrderGatewayLog.POSITION:
      assert pb.HasField('position')
      self._position[acct_key] = pb.position
    elif pb.type == OrderGatewayLog.TRANSFER:
      assert pb.HasField('transfer')
      self._transfer[acct_key].append(pb.transfer)
    elif pb.type == OrderGatewayLog.ORDER_EVENT:
      assert pb.HasField('event')
      assert pb.event.type == OrderEvent.ORDER_FILLED
      fill_key = self._fill_key_func(pb.event)
      if fill_key not in self._fill[acct_key]:
        self._fill[acct_key][fill_key] = pb.event
    else:
      pass

  def query_account_balance_history(self, request_list, query_timestamp=None):
    self._balance.clear()
    topics = _filter_account_topic(self._topics, request_list)
    if query_timestamp is not None:
      assert isinstance(query_timestamp, int)
      timestamp_from = query_timestamp - (10 * 60 * 10 ** 9)
      timestamp_to = query_timestamp
    else:
      timestamp_from = to_timestamp_int(
          datetime.datetime.utcnow() - datetime.timedelta(minutes=10))
      timestamp_to = to_timestamp_int(
          datetime.datetime.utcnow() + datetime.timedelta(minutes=10))
    run_from_kafka_archive(topics=topics,
                           kafka_config=self._kafka_config,
                           timestamp_from=timestamp_from,
                           timestamp_to=timestamp_to,
                           callback=self._on_log)

    result_list = []
    for request in request_list:
      try:
        assert request.account_id is not None
        balance = self._balance[self._acct_key_func(request)]
        balance_result = DBOutputBalance(
            account_id=request.account_id,
            query_timestamp=time.time(),
            balance=balance,
            balance_hid=None)
        result_list.append(balance_result)
      except Exception as e:
        logging.error(f'Fail to query balance {request}')
        logging.error(traceback.format_exc())
        result_list.append(None)
        continue
    return result_list

  def query_account_position_history(self, request_list, query_timestamp=None):
    self._position.clear()
    topics = _filter_account_topic(self._topics, request_list)
    if query_timestamp is not None:
      assert isinstance(query_timestamp, int)
      timestamp_from = query_timestamp - (10 * 60 * 10 ** 9)
      timestamp_to = query_timestamp
    else:
      timestamp_from = to_timestamp_int(
          datetime.datetime.utcnow() - datetime.timedelta(minutes=10))
      timestamp_to = to_timestamp_int(
          datetime.datetime.utcnow() + datetime.timedelta(minutes=10))
    run_from_kafka_archive(topics=topics,
                           kafka_config=self._kafka_config,
                           timestamp_from=timestamp_from,
                           timestamp_to=timestamp_to,
                           callback=self._on_log)

    result_list = []
    for request in request_list:
      try:
        assert request.account_id is not None
        position = self._position[self._acct_key_func(request)]
        position_result = DBOutputPosition(
            account_id=request.account_id,
            query_timestamp=time.time(),
            position=position,
            position_hid=None)
        result_list.append(position_result)
      except Exception as e:
        logging.error(f'Fail to query position {request}')
        logging.error(traceback.format_exc())
        result_list.append(None)
        continue
    return result_list

  def query_transfer_history(self, request_list, start_timestamp, end_timestamp):
    self._transfer.clear()
    topics = _filter_account_topic(self._topics, request_list)
    assert isinstance(start_timestamp, int)
    assert isinstance(end_timestamp, int)
    run_from_kafka_archive(topics=topics,
                           kafka_config=self._kafka_config,
                           timestamp_from=start_timestamp,
                           timestamp_to=end_timestamp,
                           callback=self._on_log)

    result_list = []
    for request in request_list:
      try:
        assert request.market_type is not None
        assert request.exchange is not None
        assert request.owner is not None
        transfers = self._transfer[self._acct_key_func(request)]
        transfers = [elem for transfer in transfers
          for elem in transfer.each_transfer]
        transfers = sorted(transfers, key=lambda x: x.create_timestamp)
        transfer_bundle = CurrencyTransferBundle(
            market_type=request.market_type,
            exchange=request.exchange,
            each_transfer=transfers
        )
        transfer_result = DBOutputTransfer(
            query_timestamp=time.time(),
            transfer=transfer_bundle
        )
        result_list.append(transfer_result)
      except Exception as e:
        logging.error(f'Fail to query transfer {request}')
        result_list.append(None)
        continue
    return result_list

  def query_account_fill_history(self, request_list, start_timestamp, end_timestamp):
    self._fill.clear()
    topics = _filter_account_topic(self._topics, request_list)
    assert isinstance(start_timestamp, int)
    assert isinstance(end_timestamp, int)
    run_from_kafka_archive(topics=topics,
                           kafka_config=self._kafka_config,
                           timestamp_from=start_timestamp,
                           timestamp_to=end_timestamp,
                           callback=self._on_log)

    result_list = []
    for request in request_list:
      try:
        assert request.market_type is not None
        assert request.exchange is not None
        assert request.owner is not None
        fill_list = [fill for fill in
          self._fill[self._acct_key_func(request)].values()]
        fill_list = [
          ProductFillElement( 
              symbol=fill.symbol,
              order_price=fill.order_price,
              order_qty=fill.order_qty,
              price=fill.fill_price,
              qty=fill.fill_qty,
              side=fill.order_side,
              fill_timestamp=fill.event_time,
              fill_type=fill.fill_type,
              fill_id=fill.fill_id,
              client_order_id=fill.internal_order_id,
              exchange_order_id=fill.external_order_id,
              fee=fill.fee if fill.HasField('fee') else None,
              fee_currency=fill.fee_currency \
                if fill.HasField('fee_currency') else None,
              proc_order_id=fill.proc_order_id
          ) for fill in fill_list
        ]
        fill_bundle = ProductFillBundle(
            exchange=request.exchange,
            market_type=request.market_type,
            owner=request.owner,
            each_fill=fill_list
        )
        fill_result = DBOutputAccountFill(
            query_timestamp=time.time(),
            fills=fill_bundle
        )
        result_list.append(fill_result)
      except Exception as e:
        logging.error(f'Fail to query account fill {request}')
        logging.error(traceback.format_exc())
        result_list.append(None)
        continue
    return result_list


if __name__ == "__main__":
  kafka_config = KafkaConfig.from_cmd_config(
      '../../coin_deploy/support_monitor/config/kafka_aws_config_2.json')
  querier = KafkaAccountInfoQuerier(kafka_config)
  request_list = query_accounts(
      business_units=['Coin'],
      market_types=['Futures'],
      exchanges=['Okex'],
      owners=['prestoinvt20'],
      as_proto=True)
  print(querier.query_account_balance_history(request_list))
  print(querier.query_account_position_history(request_list))
